from cn.redguest.pbase.demo.Gan import GanNetwork
import torch
import numpy
from torch.autograd import Variable
from scipy.misc.pilutil import *
import shutil
import os

from tensorflow.examples.tutorials.mnist import input_data


def main():
    data_set = input_data.read_data_sets("MNIST_data/", one_hot=True)
    epoch = 1000000
    out_dir = "out"
    gan = GanNetwork().cuda()

    real_batch_size = 1
    real_input_len = 28 * 28

    gen_batch_size = 1
    gen_rand_seed_len = 3

    if os.path.exists(out_dir):
        shutil.rmtree(out_dir)

    os.mkdir(out_dir)

    for i in range(epoch):
        next_bat = data_set.train.next_batch(real_batch_size)
        real_data = next_bat[0].reshape([-1, real_input_len])
        seed_data = numpy.random.uniform(0, 1, gen_batch_size * gen_rand_seed_len).reshape(
            [-1, gen_rand_seed_len])

        gen_data = gan.train(
            Variable(torch.FloatTensor(seed_data)).cuda(),
            Variable(torch.FloatTensor(real_data)).cuda()
        )

        if i % 100 == 0:
            print("epoch %d" % i)
            gen_dat = gen_data.cpu().data.numpy()
            gen_dat = gen_dat * 255
            gen_dat = gen_dat.astype(numpy.uint8)
            gen_dat = gen_dat.reshape([real_batch_size, 28, 28])[0].reshape([28, 28])
            imsave(out_dir + "/epoch%d.jpg" % i, gen_dat)


if __name__ == "__main__":
    main()
